package com.ruoyi.flink.wc.util;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.UUID;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.CheckpointListener;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.rocketmq.client.consumer.MQPullConsumer;
import org.apache.rocketmq.client.consumer.MQPullConsumerScheduleService;
import org.apache.rocketmq.client.consumer.PullResult;
import org.apache.rocketmq.client.consumer.PullTaskCallback;
import org.apache.rocketmq.client.consumer.PullTaskContext;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;

public class RocketMQSource extends RichParallelSourceFunction<String>
    					implements CheckpointedFunction, CheckpointListener, ResultTypeQueryable<String> {

    private static final long serialVersionUID = 1L;
    private boolean isRunning =true;
    private transient MQPullConsumerScheduleService pullConsumerScheduleService;
    private String namesrvaddr;
    private String topic;
    private String group;

    public RocketMQSource(String namesrvaddr, String topic, String group) {
        this.namesrvaddr = namesrvaddr;
        this.topic = topic;
        this.group = group;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        
        pullConsumerScheduleService = new MQPullConsumerScheduleService(group);
        pullConsumerScheduleService.getDefaultMQPullConsumer().setNamesrvAddr(namesrvaddr);
        pullConsumerScheduleService.setMessageModel(MessageModel.CLUSTERING);
        pullConsumerScheduleService.getDefaultMQPullConsumer().setInstanceName(UUID.randomUUID().toString());
        pullConsumerScheduleService.getDefaultMQPullConsumer().setPersistConsumerOffsetInterval(5000);
        pullConsumerScheduleService.getDefaultMQPullConsumer().setPollNameServerInterval(30000);
        pullConsumerScheduleService.getDefaultMQPullConsumer().setHeartbeatBrokerInterval(30000);

    }

    @Override
    public void run(SourceContext<String> context) throws Exception {
        
        final Object lock = context.getCheckpointLock();
        
        pullConsumerScheduleService.registerPullTaskCallback(topic, new PullTaskCallback() {

            @Override
            public void doPullTask(MessageQueue mq, PullTaskContext pullTaskContext) {
                try {
                	MQPullConsumer consumer = pullTaskContext.getPullConsumer();
                	
                    long offset =  consumer.fetchConsumeOffset(mq, false);
                    if (offset < 0) {
                    	offset = consumer.maxOffset(mq);
                    }

                    PullResult pullResult = consumer.pull(mq, "*", offset, 32);
                    boolean found = false;
                    switch (pullResult.getPullStatus()) {
                        case FOUND:
                            List<MessageExt> messages = pullResult.getMsgFoundList();
                            for (MessageExt msg : messages) {
                            	String data = msg.getBody() != null ? new String(msg.getBody(), StandardCharsets.UTF_8) : null;
                            	if(data!=null) {
                                    synchronized (lock) {
                                        context.collectWithTimestamp(data, msg.getBornTimestamp());
                                    }
                            	}
                            }
                            found = true;
                            break;
                        case NO_MATCHED_MSG:
                            break;
                        case NO_NEW_MSG:
                            break;
                        case OFFSET_ILLEGAL:
                            break;
                        default:
                            break;
                    }

                    synchronized (lock) {
                    	consumer.updateConsumeOffset(mq, pullResult.getNextBeginOffset());
                    }

                    if (found) {
                        pullTaskContext.setPullNextDelayTimeMillis(0);
                    } else {
                        pullTaskContext.setPullNextDelayTimeMillis(10);
                    }
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
        });

        try {
            pullConsumerScheduleService.start();
        } catch (MQClientException e) {
            throw new RuntimeException(e);
        }

        while (this.isRunning) {
            Thread.sleep(1000);
        }
    }

    @Override
    public void cancel() {
        this.isRunning = false;
        if (pullConsumerScheduleService != null) {
            pullConsumerScheduleService.shutdown();
        }
    }

    @Override
    public void close() throws Exception {
        try {
            cancel();
        } finally {
            super.close();
        }
    }

	@Override
	public TypeInformation<String> getProducedType() {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public void notifyCheckpointComplete(long checkpointId) throws Exception {
		// TODO Auto-generated method stub
		
	}

	@Override
	public void snapshotState(FunctionSnapshotContext context) throws Exception {
		// TODO Auto-generated method stub
		
	}

	@Override
	public void initializeState(FunctionInitializationContext context) throws Exception {
		// TODO Auto-generated method stub
		
	}
}